import torch

def convert_to_yolo_format(im0_shape, xyxy):
    """
    Convert bounding boxes to YOLO format.
    
    Args:
        im0_shape: Tuple[int, int] - Original image size (height, width)
        xyxy: List[Tensor] - Bounding boxes in (x1, y1, x2, y2) format
    
    Returns:
        List[str] - YOLO formatted bounding boxes
    """
    # Normalize image width and height
    h, w = im0_shape
    x1, y1, x2, y2 = xyxy
    x_center = (x1 + x2) / 2 / w
    y_center = (y1 + y2) / 2 / h
    width = (x2 - x1) / w
    height = (y2 - y1) / h
    return [x_center, y_center, width, height]

def parse_detection_results(file_path):
    """
    Parse detection results from a text file and convert them to YOLO format.
    
    Args:
        file_path: str - Path to the input file with detection results
    
    Returns:
        List[str] - List of YOLO formatted lines
    """
    yolo_lines = []
    with open(file_path, 'r') as file:
        lines = file.readlines()
    
    for line in lines:
        parts = line.strip().split(' ')
        image_id = parts[0].replace('Image:', '').strip()
        confidence = parts[2]
        xyxy = [float(x) for x in parts[4].strip('[]').split(',')]
        
        # Replace with actual image size
        im0_shape = (1080, 1920)  # Example image shape, replace with actual
        
        yolo_format = convert_to_yolo_format(im0_shape, torch.tensor(xyxy))
        yolo_line = f"{image_id} rust {confidence} " + " ".join(map(str, yolo_format))
        yolo_lines.append(yolo_line)
    
    return yolo_lines

def save_yolo_format_results(yolo_lines, output_path):
    """
    Save YOLO formatted results to a file.
    
    Args:
        yolo_lines: List[str] - List of YOLO formatted lines
        output_path: str - Path to the output file
    """
    with open(output_path, 'w') as file:
        for line in yolo_lines:
            file.write(line + "\n")

# Example usage
file_path = 'rust_detect.txt'  # Replace with your actual file path
output_path = 'yolo_labels.txt'  # Replace with your desired output file path

yolo_lines = parse_detection_results(file_path)
save_yolo_format_results(yolo_lines, output_path)

